#include "MTBinaryOperator.hpp"
#include "MTExpression.hpp"
#include "matrix_grammar.tab.hpp"
#include "Params.h"
#include <string>
#include <iostream>
#include <sstream>

namespace MathematicaTranslator
{
	
	std::string* MTBinaryOperator::generateCode(CodeGenContext& context)
	{
		myContext = &context;

		std::ostringstream ss;
		
		std::string *temp1 = expr1.generateCode(context);
		std::string *temp2 = expr2.generateCode(context);

		std::string *t1 = expr1.getType();
		std::string *t2 = this->getType();
		std::string *numberOfElements;
		std::ofstream *genFileStream = context.getStream();
		std::string uniType = "double";

		if (t1->find_first_of("[") != -1 || t2->find_first_of("[") != -1){			
			int i1 = t2->find_first_of("[");		
			int i2 = t2->find_first_of("]");
			numberOfElements = new std::string(t2->substr(i1 + 1, i2 - i1 -1));

			if (this->op.compare("+") == 0){
				bool writeFunction = false;

				if (t1->find_first_of("i") == -1) {
					if (!context.listDoubleAdd) {
						context.listDoubleAdd = true;		
						writeFunction = true;
						uniType = "double";
					}
				}
				else {
					if (!context.listIntAdd) {
						context.listIntAdd = true;
						writeFunction = true;
						uniType = "int";
					}
				}

				if (writeFunction){

					*genFileStream << "__global__ void addListKernel(" << uniType << " *c, const " << uniType << " *a, const " << uniType << " *b)" << "\n";
					*genFileStream << "{" << "\n";
					*genFileStream << "	int i = threadIdx.x;" << "\n";
					*genFileStream << "	c[i] = a[i] + b[i];" << "\n";
					*genFileStream << "}" << "\n\n";

					// ==================================================================================================

					*genFileStream << "cudaError_t addWithCuda(" << uniType << " *c, const " << uniType << " *a, const " << uniType << " *b, size_t size)" << "\n";
					*genFileStream << "{" << "\n";
					*genFileStream << "	" << uniType << " *dev_a = 0;" << "\n";
					*genFileStream << "	" << uniType << " *dev_b = 0;" << "\n";
					*genFileStream << "	" << uniType << " *dev_c = 0;" << "\n";
					*genFileStream << "	cudaError_t cudaStatus;" << "\n"  << "\n";

						// Choose which GPU to run on, change this on a multi-GPU system.
					*genFileStream << "	cudaStatus = cudaSetDevice(0);" << "\n";
					*genFileStream << "	if (cudaStatus != cudaSuccess) {" << "\n";
					*genFileStream << "		fprintf(stderr, \"cudaSetDevice failed!  Do you have a CUDA-capable GPU installed?\");" << "\n";
					*genFileStream << "		goto Error;" << "\n";
					*genFileStream << "	}" << "\n" << "\n";

						// Allocate GPU buffers for three vectors (two input, one output)    .
					*genFileStream << "	cudaStatus = cudaMalloc((void**)&dev_c, size * sizeof(" << uniType << "));" << "\n";
					*genFileStream << "	if (cudaStatus != cudaSuccess) {" << "\n";
					*genFileStream << "		fprintf(stderr, \"cudaMalloc failed!\");" << "\n";
					*genFileStream << "		goto Error;" << "\n";
					*genFileStream << "	}" << "\n" << "\n";

					*genFileStream << "	cudaStatus = cudaMalloc((void**)&dev_a, size * sizeof(" << uniType << "));" << "\n";
					*genFileStream << "	if (cudaStatus != cudaSuccess) {" << "\n";
					*genFileStream << "		fprintf(stderr, \"cudaMalloc failed!\");" << "\n";
					*genFileStream << "		goto Error;" << "\n";
					*genFileStream << "	}" << "\n" << "\n";

					*genFileStream << "	cudaStatus = cudaMalloc((void**)&dev_b, size * sizeof(" << uniType << "));" << "\n";
					*genFileStream << "	if (cudaStatus != cudaSuccess) {" << "\n";
					*genFileStream << "		fprintf(stderr, \"cudaMalloc failed!\");" << "\n";
					*genFileStream << "		goto Error;" << "\n";
					*genFileStream << "	}" << "\n" << "\n";

						// Copy input vectors from host memory to GPU buffers.
					*genFileStream << "	cudaStatus = cudaMemcpy(dev_a, a, size * sizeof(" << uniType << "), cudaMemcpyHostToDevice);" << "\n";
					*genFileStream << "	if (cudaStatus != cudaSuccess) {" << "\n";
					*genFileStream << "		fprintf(stderr, \"cudaMemcpy failed!\");" << "\n";
					*genFileStream << "		goto Error;" << "\n";
					*genFileStream << "	}" << "\n" << "\n";

					*genFileStream << "	cudaStatus = cudaMemcpy(dev_b, b, size * sizeof(" << uniType << "), cudaMemcpyHostToDevice);" << "\n";
					*genFileStream << "	if (cudaStatus != cudaSuccess) {" << "\n";
					*genFileStream << "		fprintf(stderr, \"cudaMemcpy failed!\");" << "\n";
					*genFileStream << "		goto Error;" << "\n";
					*genFileStream << "	}" << "\n" << "\n";

						// Launch a kernel on the GPU with one thread for each element.
					*genFileStream << "	addListKernel<<<1, size>>>(dev_c, dev_a, dev_b);" << "\n" << "\n";

						// cudaDeviceSynchronize waits for the kernel to finish, and returns
						// any errors encountered during the launch.
					*genFileStream << "	cudaStatus = cudaDeviceSynchronize();" << "\n";
					*genFileStream << "	if (cudaStatus != cudaSuccess) {" << "\n";
					*genFileStream << "		fprintf(stderr, \"cudaDeviceSynchronize returned error code %d after launching addListKernel!\", cudaStatus);" << "\n";
					*genFileStream << "		goto Error;" << "\n";
					*genFileStream << "	}" << "\n" << "\n";

						// Copy output vector from GPU buffer to host memory.
					*genFileStream << "	cudaStatus = cudaMemcpy(c, dev_c, size * sizeof(" << uniType << "), cudaMemcpyDeviceToHost);" << "\n";
					*genFileStream << "	if (cudaStatus != cudaSuccess) {" << "\n";
					*genFileStream << "		fprintf(stderr, \"cudaMemcpy failed!\");" << "\n";
					*genFileStream << "		goto Error;" << "\n";
					*genFileStream << "	}" << "\n" << "\n";

					*genFileStream << "	Error:" << "\n";
					*genFileStream << "		cudaFree(dev_c);" << "\n";
					*genFileStream << "		cudaFree(dev_a);" << "\n";
					*genFileStream << "		cudaFree(dev_b);" << "\n" << "\n";

					*genFileStream << "		return cudaStatus;" << "\n";
					*genFileStream << "	}" << "\n\n";
				}
			}
			else if (this->op.compare("*") == 0){
			
			}


			// writing to ss 
			//uniType = *t2;
			std::ostringstream v0, v1, v2, v3, v4;
		
			v0 << MTNode::tmpVariable << MTNode::tmpVariableCount;
			std::string var0 = v0.str(); 
			//MTNode::tmpVariableCount++;

			std::string var1 = "0"; 
			std::string var2 = "0";
			
			if (expr1.getClassType() == NORMALBINARYOP){
				ss << "const " << uniType << " " << var0 << " = { 0 };" << "\n";
				v1 << MTNode::tmpVariable << MTNode::tmpVariableCount+1;
				var1 = v1.str();
				MTNode::tmpVariableCount++;
				ss << *expr1.generateCode(context);
				var2 = *expr2.generateCode(context);
			}
			else {
				var1 = *expr1.generateCode(context);
				var2 = *expr2.generateCode(context);
			}
					
			if (op.compare("+") == 0){
				ss << "\ncudaError_t cudaStatus = addWithCuda(" << var0 << ", " << var1 << ", " << var2 << ", " << *numberOfElements << ");" << "\n";
			}
			else if (op.compare("*") == 0){
				ss << "\ncudaError_t cudaStatus = mulWithCuda(" << var0 << ", " << var1 << ", " << var2 << ", " << *numberOfElements << ");" << "\n";
			}
			ss << "if (cudaStatus != cudaSuccess) {" << "\n";
			ss << "	fprintf(stderr, \"addWithCuda failed!\");" << "\n";
			ss << "	return 1;" << "\n";
			ss << "}" << "\n";
				//ss << "lista"; // lista
		}
		else {
			ss << *expr1.generateCode(context);
			ss << " ";
			ss << op;
			ss << " ";
			ss << *expr2.generateCode(context);
		}

		return new std::string(ss.str());
	}
		
	std::string* MTBinaryOperator::getType(){
		std::string *expr1_type = expr1.getType();
		std::string *expr2_type = expr2.getType();

		if (expr1_type->compare("id") == 0){
			if (myContext->locals().find(*expr1.generateCode(*myContext)) != myContext->locals().end()){
				std::string *temp = myContext->locals().find(*expr1.generateCode(*myContext))->second;
				if (temp->find_first_of("[") != -1){
					int afterType = temp->find_first_of(" ");
					int size1 = temp->find_first_of("[");
					int size2 = temp->find_first_of("]");

					std::ostringstream ss;

					ss << temp->substr(0,afterType);
					ss << temp->substr(size1, size2 - size1 + 1);

					expr1_type = new std::string(ss.str());

					//std::cout << "==========> MAM ID W EXPR 1 z LISTA LUB MACIERZA: " << ss.str() << " #\n";
				}
				else {
					int afterType = temp->find_first_of(" ");

					std::ostringstream ss;

					ss << temp->substr(0,afterType);

					expr1_type = new std::string(ss.str());

					//std::cout << "==========> MAM ID W EXPR 1 z ZWYKLA ZMIENNA: " << ss.str() << "\n";
				}
			}
			else {
				std::cerr << "ID not found" << std::endl;
			}
		}
		if (expr2_type->compare("id") == 0){
			if (myContext->locals().find(*expr2.generateCode(*myContext)) != myContext->locals().end()){
				std::string *temp = myContext->locals().find(*expr2.generateCode(*myContext))->second;
				if (temp->find_first_of("[") != -1){
					int afterType = temp->find_first_of(" ");
					int size1 = temp->find_first_of("[");
					int size2 = temp->find_first_of("]");

					std::ostringstream ss;

					ss << temp->substr(0,afterType);
					ss << temp->substr(size1, size2 - size1 + 1);

					expr2_type = new std::string(ss.str());

					//std::cout << "==========> MAM ID W EXPR 2 z LISTA LUB MACIERZA\n";
				}
				else {
					int afterType = temp->find_first_of(" ");

					std::ostringstream ss;

					ss << temp->substr(0,afterType);

					expr2_type = new std::string(ss.str());

					//std::cout << "==========> MAM ID W EXPR 2 z ZWYKLA ZMIENNA \n";
				}
			}
			else {
				std::cerr << "ID not found" << std::endl;
			}
		}

		int k1 = expr1_type->find_first_of("[");
		int k2 = expr2_type->find_first_of("[");
		if (k1 != -1 || k2 != -1){
			if (k1 != -1 && k2 != -1){
				int k1b = expr1_type->find_first_of("]");
				int k2b = expr2_type->find_first_of("]");
				
				std::string res1 = expr1_type->substr(k1, k1b - k1 + 1);
				std::string res2 = expr2_type->substr(k2, k2b - k2 + 1);

				if (res1.compare(res2) != 0) {
					std::cerr << "DIFFRENT SIZE" << std::endl;
					return new std::string("WRONG SIZE");
				}
								
				return expr1_type;
			}
			else {
				std::cerr << "WRONG TYPE" << std::endl;
			}
		}

		std::string *type = new std::string("float");
		if (*expr1_type < *type) return expr1_type;
		if (*expr2_type < *type) return expr2_type; 
		if (*expr1_type > *type && *expr2_type > *type) return expr1_type;
		return type;
	}

	int MTBinaryOperator::getClassType(){
		return NORMALBINARYOP;
	}
}